import pynvml
import time
import smtplib
from email.mime.text import MIMEText 
from email.header import Header

threshold = 2000
mail_host = "xxx.xxx"
mail_user = "xxx@xxx.xxx"
mail_pass = "password"
sender = "name"
receiver = 'xxx@xxx.xxx'


def mail(info):
    msg = MIMEText(info, 'plain', 'utf-8')
    msg['From'] = Header(mail_user, 'utf-8')
    msg['To'] = Header(receiver, 'utf-8')

    subject = '3090 GPU 空闲！'
    msg['Subject'] = Header(subject, 'utf-8')

    server = smtplib.SMTP()
    server.connect(mail_host, 25)
    server.login(mail_user, mail_pass)
    server.sendmail(mail_user, [receiver], msg.as_string())
    server.quit()
    print("SUCCESS")



pynvml.nvmlInit()
already_sent = False
while True:
    flag = False
    usage = [0, 0]
    for num in range(2):
        handle = pynvml.nvmlDeviceGetHandleByIndex(num)
        meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
        usage[num] = int(meminfo.used/1024/1024)
        if int(meminfo.used/1024/1024) < threshold:
            flag = True
    if not flag:
        already_sent = False
    if flag and not already_sent:
        already_sent = True
        info = f'''
Server: SUN 133.9.96.59:22103
GPU Usage :
            GPU 0: {usage[0]}M
            GPU 1: {usage[1]}M
'''
        mail(info)
        time.sleep(1800)
    elif flag:
        print("Have casual GPU but already sent mail!!!!")
    else:
        print("ALL BUSY!!!!")
    
    time.sleep(120)
